from fastapi import UploadFile
from tortoise.expressions import Q

from client import CosClient
from utils import file_utils, utils, excel_utils
from models import DataPool
from models import DataPoolExt
from models import DataPoolTag
import pandas as pd
from service import BaseService

COLUMNS_DICT = excel_utils.columns_dict


class DataPoolService(BaseService):

    def __init__(self, user_id, tenant_id, e_id):
        super().__init__()
        self.tag_uid = None
        self.user_id = user_id
        self.tenant_id = tenant_id
        self.e_id = e_id

    async def upload_file(self, tag, file: UploadFile):
        self.tag_uid = tag
        # save file path
        org_file_name = file.filename
        # write data to file
        file_path = await file_utils.save_file_to_temp(file)
        # write data to remote
        remote_file_path = file_utils.get_cos_file_path(file_path)
        CosClient.upload_file(remote_file_path, file_path)
        # read file
        await excel_utils.read_file(file_path, self._format_data, self._save_data)

        file_utils.del_temp_file(file_path)

    async def _format_data(self, line, ext_list: list):
        print("---------------tag:{}".format(self.tenant_id))
        data_dict = {}
        for item in COLUMNS_DICT.keys():
            name = line.get(item, None)
            if name and str(name) != "nan":
                data_dict[COLUMNS_DICT[item]] = name
        uid = utils.get_uid()
        data_dict.update(data_source=0,
                         create_user=self.user_id,
                         update_user=self.user_id,
                         uid=uid,
                         e_id=self.e_id,
                         tenant_id=self.tenant_id)

        dp = DataPool(**data_dict)

        dp_ext_list = []
        for item in ext_list:
            name = line.get(item, None)
            if name and str(name) != "nan":
                ext_data = {"name": item, "field": item, "value": name, "pool_id": uid}
                ext_data.update(
                    create_user=self.user_id,
                    update_user=self.user_id,
                    uid=utils.get_uid(),
                    e_id=self.e_id,
                    tenant_id=self.tenant_id)
                dp_ext_list.append(DataPoolExt(**ext_data))

        return dp, dp_ext_list

    async def _save_data(self, pool_list: list, pool_ext_list: list):
        print("save data")
        if len(pool_list):
            await DataPool.bulk_create(pool_list)
            await self._save_tag(pool_list)

        if len(pool_ext_list):
            await DataPoolExt.bulk_create(pool_ext_list)

    async def _save_tag(self, pool_list):
        tag_list = [await self._parse_tag(item) for item in pool_list]

        await DataPoolTag.bulk_create(tag_list)

    async def _parse_tag(self, pool: DataPool) -> DataPoolTag:
        data_dict = {"tag_id": self.tag_uid, "pool_id": pool.uid}
        data_dict.update(
            create_user=self.user_id,
            update_user=self.user_id,
            uid=utils.get_uid(),
            e_id=self.e_id,
            tenant_id=self.tenant_id)
        return DataPoolTag(**data_dict)

    @classmethod
    async def list_data(cls):
        return await cls.page_list(DataPool.filter(Q(is_delete=0)))
